#!/usr/bin/env python3
"""
Qwen2.5-32B 快速测试脚本
用于验证模型是否能正常加载和推理
"""

import sys
import os

def check_conda_env():
    """检查是否在正确的 conda 环境中"""
    import sys
    import os
    
    # 检查是否在 conda 环境中
    conda_env = os.environ.get('CONDA_DEFAULT_ENV', '')
    if conda_env:
        print(f"  当前 Conda 环境: {conda_env}")
        if conda_env != 'qwen32b':
            print(f"  ⚠ 警告: 建议在 'qwen32b' 环境中运行")
            print(f"    请运行: conda activate qwen32b")
    else:
        print("  ⚠ 警告: 未检测到 Conda 环境")
        print("    建议运行: conda activate qwen32b")
    
    return conda_env == 'qwen32b'

def test_imports():
    """测试必要的包是否已安装"""
    print("【1/4】检查 Python 包...")
    
    # 检查 conda 环境
    in_correct_env = check_conda_env()
    
    torch_available = False
    cuda_available = False
    vllm_available = False
    
    # 检查 PyTorch
    try:
        import torch
        torch_available = True
        print(f"  ✓ PyTorch: {torch.__version__}")
        
        cuda_available = torch.cuda.is_available()
        print(f"  {'✓' if cuda_available else '✗'} CUDA 可用: {cuda_available}")
        
        if cuda_available:
            print(f"  ✓ GPU 数量: {torch.cuda.device_count()}")
            print(f"  ✓ CUDA 版本: {torch.version.cuda}")
        else:
            print("  ⚠ CUDA 不可用，可能的原因：")
            print("    1. 系统没有 GPU")
            print("    2. NVIDIA 驱动未安装")
            print("    3. CUDA 运行时未正确配置")
            print("    4. PyTorch 安装的是 CPU 版本")
            script_dir = os.path.dirname(os.path.abspath(__file__))
            print(f"    建议运行: bash {script_dir}/diagnose_gpu.sh")
    except ImportError:
        print("  ✗ PyTorch 未安装")
        print("    安装命令: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121")
    
    # 检查 vLLM
    try:
        import vllm
        vllm_available = True
        print(f"  ✓ vLLM: {vllm.__version__}")
    except ImportError:
        print("  ✗ vLLM 未安装")
        print("    安装命令: pip install vllm")
        if not in_correct_env:
            print("    注意: 请先激活 qwen32b 环境: conda activate qwen32b")
    
    # 如果关键包缺失，提供安装建议
    if not torch_available or not vllm_available:
        print("\n  💡 安装建议：")
        if not in_correct_env:
            print("    1. 激活环境: conda activate qwen32b")
        if not torch_available:
            print("    2. 安装 PyTorch: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121")
        if not vllm_available:
            print("    3. 安装 vLLM: pip install vllm")
        print("    或者重新运行部署脚本: bash deploy_qwen.sh")
        return False
    
    return True

def test_model_files():
    """检查模型文件是否存在"""
    print("\n【2/4】检查模型文件...")
    model_path = "./models/Qwen/Qwen2.5-32B"
    
    if not os.path.exists(model_path):
        print(f"  ✗ 模型目录不存在: {model_path}")
        return False
    
    print(f"  ✓ 模型目录存在: {model_path}")
    
    required_files = ["config.json", "tokenizer_config.json"]
    for file in required_files:
        file_path = os.path.join(model_path, file)
        if os.path.exists(file_path):
            print(f"  ✓ {file} 存在")
        else:
            print(f"  ⚠ {file} 不存在（可能不影响使用）")
    
    return True

def check_cuda_available():
    """检查 CUDA 是否可用"""
    try:
        import torch
        if torch.cuda.is_available():
            return True, torch.cuda.device_count()
        else:
            return False, 0
    except:
        return False, 0

def test_model_loading():
    """测试模型加载"""
    print("\n【3/4】测试模型加载...")
    model_path = "./models/Qwen/Qwen2.5-32B"
    
    # 首先检查 CUDA 是否可用
    cuda_available, gpu_count = check_cuda_available()
    
    if not cuda_available:
        print("  ⚠ 警告: 未检测到可用的 CUDA/GPU")
        print("  vLLM 需要 GPU 支持才能运行。")
        print("\n  可能的原因：")
        print("    1. 系统没有 GPU")
        print("    2. NVIDIA 驱动未安装或版本过旧")
        print("    3. CUDA 工具包未正确安装")
        print("    4. PyTorch 未安装 CUDA 版本")
        print("\n  检查建议：")
        print("    - 运行: nvidia-smi 查看 GPU 状态")
        print("    - 检查 PyTorch CUDA 版本: python -c 'import torch; print(torch.cuda.is_available())'")
        print("    - 如果确实没有 GPU，Qwen2.5-32B 模型太大，无法在 CPU 上运行")
        print("\n  注意: Qwen2.5-32B 需要至少 80GB GPU 内存（如 A100 80G）")
        
        # 询问是否尝试使用 transformers 进行 CPU 测试（仅用于验证模型文件）
        response = input("\n  是否尝试使用 transformers 库验证模型文件（仅检查，不进行推理）? [y/N]: ")
        if response.lower() in ['y', 'yes']:
            return test_model_with_transformers(model_path)
        else:
            return None
    
    print(f"  ✓ 检测到 {gpu_count} 个 GPU，开始加载模型...")
    
    try:
        from vllm import LLM
        
        print("  正在加载模型（这可能需要几分钟，请耐心等待）...")
        # 使用较小的 max_model_len 以加快测试速度
        llm = LLM(
            model=model_path,
            trust_remote_code=True,
            max_model_len=2048,  # 可以根据 GPU 内存调整
            gpu_memory_utilization=0.8  # 使用 80% 的 GPU 内存
        )
        print("  ✓ 模型加载成功！")
        return llm
    except RuntimeError as e:
        error_msg = str(e)
        if "Device string must not be empty" in error_msg or "No platform detected" in error_msg:
            print(f"  ✗ 模型加载失败: 无法检测到 GPU 设备")
            print("  这通常意味着：")
            print("    1. vLLM 无法检测到 CUDA 设备")
            print("    2. 可能需要检查 CUDA 环境变量")
            print("    3. 或者 GPU 驱动/运行时有问题")
            print("\n  建议检查：")
            print("    - 运行: nvidia-smi")
            print("    - 检查: echo $CUDA_VISIBLE_DEVICES")
            print("    - 验证: python -c 'import torch; print(torch.cuda.is_available(), torch.cuda.device_count())'")
        else:
            print(f"  ✗ 模型加载失败: {e}")
        import traceback
        traceback.print_exc()
        return None
    except Exception as e:
        print(f"  ✗ 模型加载失败: {e}")
        import traceback
        traceback.print_exc()
        return None

def test_model_with_transformers(model_path):
    """使用 transformers 库验证模型文件（仅检查，不进行推理）"""
    print("\n  使用 transformers 库验证模型文件...")
    try:
        from transformers import AutoConfig, AutoTokenizer
        
        print("  正在加载模型配置...")
        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
        print(f"  ✓ 模型配置加载成功")
        print(f"    模型类型: {config.model_type}")
        print(f"    词汇表大小: {config.vocab_size}")
        
        print("  正在加载分词器...")
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        print(f"  ✓ 分词器加载成功")
        
        print("\n  ✓ 模型文件验证通过（使用 transformers）")
        print("  注意: 这只是验证模型文件完整性，实际推理需要 GPU 和 vLLM")
        return "transformers_verified"  # 返回特殊标记
    except ImportError:
        print("  ✗ transformers 库未安装")
        print("  可以运行: pip install transformers")
        return None
    except Exception as e:
        print(f"  ✗ 模型验证失败: {e}")
        import traceback
        traceback.print_exc()
        return None

def test_inference(llm):
    """测试推理"""
    print("\n【4/4】测试推理...")
    
    if llm is None:
        print("  ✗ 跳过推理测试（模型未加载）")
        return False
    
    if llm == "transformers_verified":
        print("  ℹ 使用 transformers 验证模式，跳过推理测试")
        print("  要进行实际推理，需要：")
        print("    1. 可用的 GPU（推荐 A100 80G 或更高）")
        print("    2. 正确配置的 CUDA 环境")
        print("    3. vLLM 能够检测到 GPU")
        return True
    
    try:
        from vllm import SamplingParams
        
        prompts = [
            "你好，请用一句话介绍一下你自己。",
        ]
        
        sampling_params = SamplingParams(
            temperature=0.7,
            top_p=0.9,
            max_tokens=100
        )
        
        print("  正在生成回复...")
        outputs = llm.generate(prompts, sampling_params)
        
        for output in outputs:
            prompt = output.prompt
            generated_text = output.outputs[0].text
            print(f"\n  提示: {prompt}")
            print(f"  回复: {generated_text}")
        
        print("\n  ✓ 推理测试成功！")
        return True
    except Exception as e:
        print(f"  ✗ 推理测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def main():
    print("=" * 50)
    print("Qwen2.5-32B 快速测试")
    print("=" * 50)
    
    # 检查包
    imports_ok = test_imports()
    
    if not imports_ok:
        print("\n" + "=" * 50)
        print("环境检查未通过")
        print("=" * 50)
        print("\n请按照上述建议安装缺失的依赖包。")
        print("\n快速安装方法：")
        script_dir = os.path.dirname(os.path.abspath(__file__))
        print(f"  bash {script_dir}/install_missing_deps.sh")
        print("\n或手动安装：")
        print("  1. 确保在正确的 conda 环境中: conda activate qwen32b")
        print("  2. 安装缺失的包（见上方提示）")
        print("\n如果已安装但仍有问题，请检查：")
        print(f"  - 运行诊断脚本: bash {script_dir}/diagnose_gpu.sh")
        print(f"  - 重新运行部署脚本: bash {script_dir}/deploy_qwen.sh")
        
        # 询问是否继续（即使有缺失的包，也可以检查模型文件）
        response = input("\n是否继续检查模型文件？(y/N): ")
        if response.lower() not in ['y', 'yes']:
            sys.exit(1)
    
    # 检查模型文件
    if not test_model_files():
        print("\n✗ 模型文件检查失败")
        sys.exit(1)
    
    # 检查 CUDA 可用性并给出提示
    cuda_available, gpu_count = check_cuda_available()
    print("\n" + "=" * 50)
    if cuda_available:
        print(f"检测到 {gpu_count} 个 GPU，可以进行完整测试")
    else:
        print("⚠ 警告: 未检测到 GPU，vLLM 无法运行")
        print("  将只能进行基础检查和模型文件验证")
    
    response = input("\n是否进行完整的模型加载和推理测试？(需要 GPU，可能需要几分钟) [y/N]: ")
    
    if response.lower() not in ['y', 'yes']:
        print("\n✓ 基础检查通过，跳过完整测试")
        print("\n提示: 运行完整测试请执行: python quick_test_qwen.py")
        sys.exit(0)
    
    # 加载模型
    llm = test_model_loading()
    
    # 推理测试
    if test_inference(llm):
        print("\n" + "=" * 50)
        print("✓ 所有测试通过！Qwen2.5-32B 已准备就绪。")
        print("=" * 50)
        sys.exit(0)
    else:
        print("\n" + "=" * 50)
        print("✗ 测试失败，请检查错误信息")
        print("=" * 50)
        sys.exit(1)

if __name__ == "__main__":
    main()

